from __future__ import print_function, annotations

from llvmlite import ir

from common.utils import log_and_raise
from spel4llvm.err_defs import ParseError
from spel4llvm.grammar_value_types import (
    TP_INT,
    TP_FLOAT,
    TP_BOOL,
    TP_STR,
    GrammarValueType,
    TP_NULL,
)
from spel4llvm.llvm_types import (
    LLVM_INT_TYPE,
    LLVM_FLOAT_TYPE,
    LLVM_BOOL_TYPE,
    LLVM_NULL_TYPE,
)


class GrammarNode:
    def __init__(self, value_type):
        self.value_type = value_type

    def to_llvm(self, builder, gen_ctx):
        pass

    def set_value_type(self, value_type):
        self.value_type = value_type

    # convert grammar node to llvm-ir
    def load(self, builder: ir.IRBuilder, gen_ctx):
        if self.value_type.is_list_type():
            # to_llvm of list is {i32, etype*}*, so need to load
            val = builder.load(self.to_llvm(builder, gen_ctx))
        else:
            # to_llvm of basic type is i32/double/i8*/i1, no need to load
            val = self.to_llvm(builder, gen_ctx)
        return val


class GrammarNodeConstant(GrammarNode):
    def __init__(self, value, value_type):
        super().__init__(value_type)
        self.value = value

    # convert to llvm value of basic type, except str
    def to_llvm(self, builder, gen_ctx):
        if self.value_type == TP_INT:
            return LLVM_INT_TYPE(self.value)
        elif self.value_type == TP_FLOAT:
            return LLVM_FLOAT_TYPE(self.value)
        elif self.value_type == TP_BOOL:
            return LLVM_BOOL_TYPE(self.value)
        elif self.value_type == TP_STR:
            # str is stored as i8**, so we load it as i8*
            return builder.load(gen_ctx.get_const_str(self.value))
        elif self.value_type == TP_NULL:
            return LLVM_NULL_TYPE
        else:
            log_and_raise(TypeError, "to_llvm unsupported in %s" % self)

    def __repr__(self):
        return "GrammarNodeConstant(value=%s, type=%s)" % (self.value, self.value_type)


GrammarNode_True = GrammarNodeConstant(1, TP_BOOL)

GrammarNode_False = GrammarNodeConstant(0, TP_BOOL)

GrammarNode_Null = GrammarNodeConstant(None, TP_NULL)


class GrammarNodeTempVar(GrammarNode):
    def __init__(self, id_name, value_type):
        super().__init__(value_type)
        self.id = id_name

    # convert to llvm's local variable returned by a llvm instruction
    def to_llvm(self, builder, gen_ctx):
        tmp_var = gen_ctx.get_tmp_var(self.id)
        if tmp_var is None:
            log_and_raise(ParseError, "temp var:%s NOT defined" % self.id)
        return tmp_var

    def __repr__(self):
        return "GrammarNodeTempVar(id=%s, type=%s)" % (self.id, self.value_type)


class GrammarNodeGlobalVar(GrammarNode):
    def __init__(self, id_name, value_type, user_visible=False):
        super().__init__(value_type)
        self.id = id_name
        self.user_visible = user_visible

    def to_llvm(self, builder, gen_ctx):
        gvar = gen_ctx.get_global_var(self.id)
        if gvar is None:
            log_and_raise(TypeError, "global var:%s NOT collected ahead" % self.id)

        return self.load_llvm_globalvar(gvar, self.value_type, builder)

    # if basic-type stored in global-var, we should call load before use
    # list-type is used as {i32, etype*}*, so it doesn't need load
    def load_llvm_globalvar(
        self,
        gvar: ir.GlobalVariable,
        value_type: GrammarValueType,
        builder: ir.IRBuilder,
    ):
        if value_type.is_basic_type():
            return builder.load(gvar)
        else:
            return gvar

    def __repr__(self):
        return "GrammarNodeGlobalVar(id=%s, type=%s, user_visible=%s)" % (
            self.id,
            self.value_type,
            self.user_visible,
        )


def can_initialize(node: GrammarNode):
    # str is i8**, cannot set to llvm global-var initializer
    return isinstance(node, GrammarNodeConstant) and node.value_type in (
        TP_INT,
        TP_FLOAT,
        TP_BOOL,
    )


# ir.Constant need raw python value when int/float/boolean, NOT llvm value
def get_llvm_constant_value(node: GrammarNode, builder, gen_ctx):
    if can_initialize(node):
        return node.value
    else:
        return node.to_llvm(builder, gen_ctx)
